schema/type.py (1)

4-8: Consider enhancing the enum with docstring, type hints, and additional data types.

The enum could benefit from:

1. Class-level documentation explaining its purpose and usage
2. Additional common data types that might be needed for metadata (e.g., DATETIME, ARRAY, OBJECT)
3. Type hints for better static analysis

Consider applying these enhancements:

 from enum import Enum
+from typing import Any
 
 
-class DataType(Enum):
+class DataType(str, Enum):
+    """Enumeration of supported data types for metadata fields.
+    
+    This enum is used for specifying metadata date formats and filtering in query engines.
+    Inheriting from str allows for direct string comparison without accessing .value
+    """
     INTEGER = "INTEGER"
     STRING = "STRING"
     BOOLEAN = "BOOLEAN"
     FLOAT = "FLOAT"
+    DATETIME = "DATETIME"
+    ARRAY = "ARRAY"
+    OBJECT = "OBJECT"
+
+    def validate_value(self, value: Any) -> bool:
+        """Validate if a value matches the declared type."""
+        type_validators = {
+            self.INTEGER: lambda x: isinstance(x, int),
+            self.FLOAT: lambda x: isinstance(x, float),
+            self.STRING: lambda x: isinstance(x, str),
+            self.BOOLEAN: lambda x: isinstance(x, bool),
+            self.DATETIME: lambda x: isinstance(x, (str, int)) and self._is_valid_datetime(x),
+            self.ARRAY: lambda x: isinstance(x, (list, tuple)),
+            self.OBJECT: lambda x: isinstance(x, dict)
+        }
+        return type_validators[self](value)

subquery.py (1)

129-138: Consider adding logging for engine selection

Adding debug logging would help track which engine is being used and why, making it easier to diagnose issues in production.

+ import logging
+ 
+ logger = logging.getLogger(__name__)
+
  if telegram and check_collection("telegram"):
      # checking if the summaries was available
      if check_collection("telegram_summary"):
+         logger.debug("Using TelegramDualQueryEngine with summaries")
          telegram_query_engine = TelegramDualQueryEngine(
              community_id=community_id
          ).prepare()
      else:
+         logger.debug("Falling back to TelegramQueryEngine (no summaries available)")
          telegram_query_engine = TelegramQueryEngine(
              community_id=community_id
          ).prepare()

utils/query_engine/base_qdrant_engine.py (2)

25-33: Consider adding error handling for setup_engine

The DualQdrantRetrievalEngine.setup_engine method may raise exceptions if initialization fails. Consider adding error handling to manage potential exceptions and ensure robustness.

---

25-33: Consider adding a docstring to the prepare method

Including a docstring for the prepare method would enhance code readability and maintainability by providing clear documentation of its purpose and usage.

utils/query_engine/qdrant_query_engine_utils.py (2)

90-91: Correct the docstring to reference the proper metadata key

In the docstring for raw_nodes, self.me seems incorrect. It should refer to self.metadata_date_key to accurately describe the metadata field used.

Apply this diff to correct the docstring:

    raw_nodes : list[NodeWithScore]
-       list of raw nodes containing metadata with self.me as float timestamp
+       list of raw nodes containing metadata with `self.metadata_date_key` as a float timestamp
        and 'text' field

---

116-120: Preserve summary order when removing duplicates

Using a set to remove duplicates from summaries may alter the original order, which could be important for readability or contextual meaning. To preserve order while removing duplicates, consider using an ordered approach.

Apply this diff to maintain the order:

- summaries = summary_text.split("\n")
- summary_bullets = set(summaries)
+ summary_bullets = []
+ for line in summary_text.split("\n"):
+     if line and line not in summary_bullets:
+         summary_bullets.append(line)
  if "" in summary_bullets:
      summary_bullets.remove("")

utils/query_engine/dual_qdrant_retrieval_engine.py (4)

53-53: Avoid shadowing built-in function filter

The variable name filter shadows the built-in Python function filter(). This can lead to unexpected behavior and reduced code readability. Consider renaming the variable to something like qdrant_filter or filters.

Apply this diff to rename the variable:

-            filter = utils.define_raw_data_filters(dates=dates)
+            qdrant_filter = utils.define_raw_data_filters(dates=dates)

And update its usage accordingly.

---

15-24: Encapsulate qa_prompt as an instance variable

The qa_prompt is defined at the module level but used within the class methods. For better encapsulation and flexibility, consider setting qa_prompt as an instance variable initialized during object creation.

Apply this diff to move qa_prompt into the class:

-qa_prompt = PromptTemplate(
+class DualQdrantRetrievalEngine(CustomQueryEngine):
+
+    def __init__(self, ..., qa_prompt=None, ...):
+        if qa_prompt is None:
+            qa_prompt = PromptTemplate(
                 "Context information is below.\n"
                 "---------------------\n"
                 "{context_str}\n"
                 "---------------------\n"
                 "Given the context information and not prior knowledge, "
                 "answer the query.\n"
                 "Query: {query_str}\n"
                 "Answer: "
             )
+        self.qa_prompt = qa_prompt

This change allows each instance to have its own prompt template.

---

95-96: Ensure consistent collection naming conventions

In both setup_engine and setup_engine_with_summaries, the collection_name is constructed using f"{community_id}_{platform_name}". To prevent unintended behavior due to inconsistent naming, make sure community_id and platform_name are consistently formatted and validated.

Also applies to: 154-155

---

201-203: Add error handling when loading the index

When loading the vector store index, there's a possibility of encountering exceptions (e.g., if the collection doesn't exist). Adding error handling will make the code more robust and provide clearer error messages.

Apply this diff to include error handling:

     qdrant_vector = QDrantVectorAccess(collection_name=collection_name)
-    index = qdrant_vector.load_index()
+    try:
+        index = qdrant_vector.load_index()
+    except Exception as e:
+        raise RuntimeError(f"Failed to load index for collection '{collection_name}': {e}")
     return index

This catches exceptions during index loading and raises a descriptive error.

Configuration used: CodeRabbit UI
Review profile: CHILL

#!/bin/bash
# Description: Check how DataType enum is used in query engines
# Expected results: Find usage in setup_engine_with_summaries and define_raw_data_filters

# Search for DataType usage in query engine files
echo "Searching for DataType usage in query engines..."
rg -A 5 "DataType" "utils/query_engine/"

# Search for metadata date format handling
echo "Searching for metadata date format handling..."
rg -A 5 "metadata_date_format" "utils/query_engine/"

#!/bin/bash
# Description: Verify Telegram engine implementations and their integration
# Expected: Find both class definitions and their relationships

# Check the relationship between Telegram engines and DualQdrantRetrievalEngine
ast-grep --pattern 'class TelegramDualQueryEngine {
  $$$
  prepare($_) {
    $$$
  }
  $$$
}'

# Look for potential breaking changes in the original TelegramQueryEngine
ast-grep --pattern 'class TelegramQueryEngine {
  $$$
  prepare($_) {
    $$$
  }
  $$$
}'

# Check for any references to these engines in test files
rg "TelegramDualQueryEngine|TelegramQueryEngine" -t python "tests/"

#!/bin/bash
# Description: Verify the DualQdrantRetrievalEngine implementation
# Expected: Find the class definition and its core methods

# Look for the class definition and its structure
ast-grep --pattern 'class DualQdrantRetrievalEngine {
  $$$
  setup_engine_with_summaries($_) {
    $$$
  }
  $$$
}'

# Check for any TODO or FIXME comments that might indicate incomplete implementation
rg -i "TODO|FIXME|XXX|HACK|BUG" "dual_qdrant_retrieval_engine.py"

#!/bin/bash
# Let's try a different approach to verify the implementation

# First, let's find the actual file
fd "dual_qdrant.*\.py$"

# Then let's examine its content with Python-specific pattern
rg "class DualQdrantRetrievalEngine" -A 5

# Look for any method definitions within this class
rg "^[[:space:]]*def" --type py dual_qdrant

# Check imports to understand dependencies
rg "^from|^import" --type py dual_qdrant

#!/bin/bash
# Let's examine the complete implementation using the full path
rg "^[[:space:]]*def" utils/query_engine/dual_qdrant_retrieval_engine.py

# Check imports
rg "^from|^import" utils/query_engine/dual_qdrant_retrieval_engine.py

# Look for any usage of Qdrant client or configuration
rg "Qdrant|qdrant" utils/query_engine/dual_qdrant_retrieval_engine.py

utils/query_engine/dual_qdrant_retrieval_engine.py (2)

14-23: Consider moving the prompt template into the class

The qa_prompt template should be moved into the class as a class constant to improve encapsulation and make it easier to modify or extend in subclasses.

 class DualQdrantRetrievalEngine(CustomQueryEngine):
+    DEFAULT_PROMPT = PromptTemplate(
+        "Context information is below.\n"
+        "---------------------\n"
+        "{context_str}\n"
+        "---------------------\n"
+        "Given the context information and not prior knowledge, "
+        "answer the query.\n"
+        "Query: {query_str}\n"
+        "Answer: "
+    )

---

132-135: Fix docstring parameter mismatch

The docstring mentions a use_summary parameter that doesn't exist in the method signature.

-        use_summary : bool
-            whether to use the summary data or not
-            note: the summary data should be available before
-            for this option to be enabled

Configuration used: CodeRabbit UI
Review profile: CHILL

utils/query_engine/dual_qdrant_retrieval_engine.py (3)

14-23: Add type hints to the prompt template

Consider adding type hints to improve IDE support and documentation:

-qa_prompt = PromptTemplate(
+qa_prompt: PromptTemplate = PromptTemplate(

---

104-107: Remove or update outdated parameter documentation

The use_summary parameter is documented but not present in the method signature. Either remove the documentation or update the method signature if the parameter is needed.

---

163-170: Add return type documentation

Add the return type information to the docstring:

        """
        prepare the vector_store for querying data

        Parameters
        ------------
        collection_name : str
            to override the default collection_name
+            
+        Returns
+        -------
+        VectorStoreIndex
+            The loaded vector store index
        """

Configuration used: CodeRabbit UI
Review profile: CHILL